#pragma once
#include <opencv2/dnn/dnn.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>

#include <chrono>
#include <cmath>
#include <exception>
#include <fstream>
#include <iostream>
#include <limits>
#include <numeric>
#include <string>
#include <queue>
#include <vector>
#include "base_util/utils.h"

#include "common/model_config.h"
#include "tensorrt/base/trt_predictor.h"
#include "tensorrt/trtmodel_manager.h"

namespace ai {

class TRTDetPredictor : public TRTPredictor {
public:
  TRTDetPredictor(TRTPackPredictor* model, TRTModelManager* mng, LogInfo *lg);
  ~TRTDetPredictor();


private:
  int RunDet(stream::ImageBlob* blob, std::vector<BaseInfo*>& det_infos) ;

    // Preprocess the input
    std::vector<std::vector<cv::cuda::GpuMat>> preprocess(const cv::cuda::GpuMat &gpuImg);

    // Detect the objects in the image
    int detectObjects(const cv::Mat &inputImageBGR, std::vector<BaseInfo*>& det_infos);
    int detectObjects(const cv::cuda::GpuMat &inputImageBGR, std::vector<BaseInfo*>& det_infos);


    // Postprocess the output
    int postprocessDetect(std::vector<float> &featureVector, std::vector<BaseInfo*>& det_infos);

    // Postprocess the output for pose model
    int postprocessPose(std::vector<float> &featureVector, std::vector<BaseInfo*>& det_infos);

    // Postprocess the output for segmentation model
    int postProcessSegmentation(std::vector<std::vector<float>> &featureVectors, std::vector<BaseInfo*>& det_infos);


    // Draw the object bounding boxes and labels on the image
    void drawObjectLabels(cv::Mat &image, std::vector<BaseInfo*>& det_infos, unsigned int scale = 2);
    cv::cuda::GpuMat resizeKeepAspectRatioPadRightBottom(const cv::cuda::GpuMat &input, size_t height, size_t width, const cv::Scalar &bgcolor);
    cv::cuda::GpuMat blobFromGpuMats(const std::vector<cv::cuda::GpuMat> &batchInput, const std::array<float, 3> &subVals, const std::array<float, 3> &divVals, bool normalize);


private:
  std::vector<float> MeshGrid;
  int HeadNum;
  int TopK = 50;


    float m_ratio = 1;
    float m_imgWidth = 0;
    float m_imgHeight = 0;

    const int TOP_K = 100;



    // Pose estimation constant
    const int NUM_KPS = 17;
    const float KPS_THRESHOLD = 0.5f;


    // Segmentation constants
    const int SEG_CHANNELS = 32;
    const int SEG_H = 160;
    const int SEG_W = 160;
    const float SEGMENTATION_THRESHOLD = 0.5f;

    std::array<float, 3> m_subVals = {0.f, 0.f, 0.f};
    std::array<float, 3> m_divVals = {1.f, 1.f, 1.f};
    bool m_normalize = true;

};


}